import json
import os
from typing import Any, Dict, Generator, List, Tuple

import datasets

from sexpdata import loads, dumps

_CITATION = ""

_DESCRIPTION = """\
The ds is used for LoRA_IE project.
"""

_HOMEPAGE = ""

_LICENSE = ""

_TRAIN_PATH = "train.json"
_TEST_PATH = "test.json"
_VAL_PATH = "dev.json"


def get_ner_prompt(text):
    return f'''Extract all named entities in the text surrounded by three &.
Output in s expression format of (:entities ((:text "xxx" :category "xxx")...))
&&& {text} &&&
'''


def get_ee_prompt(text):
    return f'''Extract all named entities and events happened in the text surrounded by three &. Output in s 
expression format of (:entities ((:text "xxx" :category "xxx")...) :events ((:trigger "xxx" :type "xxx" :arguments ((:name "xxx" :role "xxx"))))) 
&&& {text} &&&
    '''


def get_re_prompt(text):
    return f'''Extract all named entities and relationships between these entities in the text surrounded by three &.
Output in s expression format of (:entities ((:text "xxx" :category "xxx")...) :relationships ((:head "xxx" :tail "xxx" :rel_type "xxx")...))
&&& {text} &&&
    '''


# TODO: Name of the dataset usually matches the script name with CamelCase instead of snake_case
class MyDataset(datasets.GeneratorBasedBuilder):
    """TODO: Short description of my dataset."""

    VERSION = datasets.Version("0.0.0")

    BUILDER_CONFIGS = [
        datasets.BuilderConfig(name="ner", version=VERSION, description="For NER ds script"),
        datasets.BuilderConfig(name="re", version=VERSION, description="For RE ds script"),
        datasets.BuilderConfig(name="ee", version=VERSION, description="For EE DS Script")
    ]
    DEFAULT_CONFIG_NAME = "ee"

    def _info(self) -> datasets.DatasetInfo:
        features = datasets.Features({
            "instruction": datasets.Value("string"),
            "input": datasets.Value("string"),
            "output": datasets.Value("string"),
            "history": datasets.Sequence(datasets.Sequence(datasets.Value("string")))
        })
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=features,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION
        )

    def _split_generators(self, dl_manager):
        # urls = _URLS[self.config.name]
        # data_dir = dl_manager.download_and_extract(_URL)
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    "filepath": dl_manager.download_and_extract(_TRAIN_PATH),
                    "split": "train",
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.VALIDATION,
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    "filepath": dl_manager.download_and_extract(_VAL_PATH),
                    "split": "dev",
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.TEST,
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    "filepath": dl_manager.download_and_extract(_TEST_PATH),
                    "split": "test"
                },
            ),
        ]

    # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`
    def _generate_examples(self, filepath, split):
        # TODO: This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
        # The `key` is for legacy reasons (tfds) and is not important in itself, but must be unique for each example.
        # example_dataset = json.load(open(filepath, "r", encoding="utf-8"))
        with open(filepath, "r", encoding="utf-8") as f:
            file = json.load(f)
            for key, data in enumerate(file):
                # print(data)
                if self.config.name == "ner":
                    # Yields examples as (key, example) tuples
                    # print("AN NER")
                    # print(get_ner_prompt(data["sentence"]), data["entities"])
                    yield key, {
                        "instruction": str(get_ner_prompt(data["sentence"])),
                        "input": "",
                        "output": dumps({'entities': [{'text': item["e_span"], 'category': item["e_type"]} for item in
                                                    data["entities"]]}),
                        "history": []
                    }
                elif self.config.name == "ee":
                    yield key, {
                        "instruction": get_ee_prompt(data["sentence"]),
                        "input": "",
                        "output": dumps({'entities': [{'text': item["e_span"],
                                                     'category': item["e_type"]} for item in data["entities"]],
                                       'events': [{'trigger': item["trigger"],
                                                   'type': item["type"],
                                                   'arguments': [{'name': parm["name"],
                                                                  'role': parm["role"]} for parm in item["arguments"]]}
                                                  for item in data["events"]]}),
                        "history": []
                    }
                else:
                    yield key, {
                        "instruction": get_re_prompt(data["sentence"]),
                        "input": "",
                        "output": dumps({'entities': [{'text': item["e_span"],
                                                     'category': item["e_type"]} for item in data["entities"]],
                                       'relationships': [{'head': item["head"],
                                                          'tail': item["tail"],
                                                          'rel_type': item["rel_type"]}
                                                         for item in data["relationships"]]}),
                        "history": []
                    }
